include("flux.jl")
include("geometry.jl")
include("boundary.jl")

function RHS(Q, Op, Geo, γ, t, SimType)
    """
    Purpose: Evaluate RHS in 2D Euler equations, discretized on weak form
    with a numerical flux
    """
    # set basic informations
    K = Geo.K                   # number of elements
    Np = Op.Np                  # number of soulution points at each element
    Nfp = Geo.Nfp               # number of flux points at each face
    Nfaces = Geo.NFaces         # number of faces at each element

    # solution points in interior and exterior elements
    vmapM = reshape(Geo.vmapM, Nfp*Nfaces, K)
    vmapP = reshape(Geo.vmapP, Nfp*Nfaces, K)

    # compute volume contributions
    F, G, rho, u, v, p = EulerFluxes(Q, γ)

    # compute weak derivatives
    rhsQ = zeros(Float64, size(Q))

    # calculate divergence of the discontinuous flux part
    for n=1:4
        @inbounds dFdr = Op.Drw*F[:,:,n]
        @inbounds dFds = Op.Dsw*F[:,:,n]
        @inbounds dGdr = Op.Drw*G[:,:,n]
        @inbounds dGds = Op.Dsw*G[:,:,n]
        @inbounds rhsQ[:,:,n] = (Geo.rx .* dFdr + Geo.sx .* dFds) + (Geo.ry .* dGdr + Geo.sy .* dGds)
    end

    # evaluate '-' and '+' traces of conservative variables
    QM = zeros(Float64, Nfp*Nfaces, K, 4)
    QP = zeros(Float64, Nfp*Nfaces, K, 4)

    for n=1:4
        @inbounds Qn = Q[:,:,n]
        @inbounds QM[:,:,n] = Qn[vmapM]
        @inbounds QP[:,:,n] = Qn[vmapP]
    end

    # set boundary conditions by modifying positive traces
    QP = BoundaryCondition(QP, Op, Geo, t, SimType)

    # evaluate primitive variables & flux functions at '-' and '+' traces
    fM,gM,rhoM,uM,vM,pM = EulerFluxes(QM, γ)
    fP,gP,rhoP,uP,vP,pP = EulerFluxes(QP, γ)

    # compute the Rusonov numerical fluxes
    λ = max.(sqrt.(uM.^2+vM.^2) + sqrt.(abs.(γ*pM./rhoM)),
                  sqrt.(uP.^2+vP.^2) + sqrt.(abs.(γ*pP./rhoP)))
    λ = reshape(λ, Nfp, Nfaces*K)
    λ = ones(Nfp, 1)*maximum(λ, dims = 1)
    λ = reshape(λ, Nfp*Nfaces, K)

    # Lifting fluxes
    nx = Geo.Normal[:,:,1]; ny = Geo.Normal[:,:,2]

    for n=1:4
      numerical_flux = nx.*(fP[:,:,n] + fM[:,:,n]) + ny.*(gP[:,:,n] + gM[:,:,n]) +
                       λ.*(QM[:,:,n] - QP[:,:,n])
      rhsQ[:,:,n] = rhsQ[:,:,n] - Geo.LIFT*(Geo.Fscale.*numerical_flux/2)
    end
    rhsQ
end
